//
// Created by whitby on 7/25/23.
//

#include "Utils.h"

#include <opencv2/opencv.hpp>
#include <filesystem>

using namespace cv;
using namespace std;

int main(int argc, char *argv[]) {
    if (argc != 4) {
        cout << "usage: ./convert_video_imu_gravity setting.yaml source_folder output_folder" << endl;
        return -1;
    }

    // 1. load camera parameters from yaml
    cv::FileStorage fsCamera(argv[1], cv::FileStorage::READ);
    if (!fsCamera.isOpened()) {
        cerr << "Failed to load camera parameters from " << argv[1] << endl;
        return -1;
    }
    cv::Mat K, distort, rectK;
    int rectW, rectH;
    fsCamera["CameraMatrix"] >> K;
    fsCamera["Distortion"] >> distort;
    fsCamera["RectCameraMatrix"] >> rectK;
    fsCamera["RectWidth"] >> rectW, fsCamera["RectHeight"] >> rectH;

    // 2. create a video capture object, in this case we are reading the video from a file
    const string dataFolder = argv[2];
    VideoCapture videoCapture(dataFolder + "/video.mp4");
    if (!videoCapture.isOpened()) {
        cerr << "Error opening video file" << endl;
        return -1;
    } else {
        int fps = cvRound(videoCapture.get(CAP_PROP_FPS));
        int frameCount = (int) videoCapture.get(CAP_PROP_FRAME_COUNT);
        int width = (int) videoCapture.get(CAP_PROP_FRAME_WIDTH);
        int height = (int) videoCapture.get(CAP_PROP_FRAME_HEIGHT);
        cout << "video fps: " << fps << endl;
        cout << "the number of frames: " << frameCount << endl;
        cout << "frame size: " << width << " x " << height << endl;
    }

    // 3. load camera timestamps, gravity and imu data
    vector<double> cameraTimeVec;
    loadTimeFile(dataFolder + "/times.txt", cameraTimeVec);
    int num_times = (int) cameraTimeVec.size();
    cout << "load " << num_times << " camera timestamps" << endl;

    vector<Gravity> gravityVec;
    loadGravityFile(dataFolder + "/gravity.txt", gravityVec);
    int num_gravity = (int) gravityVec.size();
    cout << "load " << num_gravity << " gravity data" << endl;

    vector<IMU> imuVec;
    loadIMUFile(dataFolder + "/imu.txt", imuVec);
    int num_imu = (int) imuVec.size();
    cout << "load " << num_imu << " imu data" << endl;

    // 4. output data
    const string outFolder = argv[3];
    const string outImgFolder = outFolder + "/images";
    if (!filesystem::exists(outImgFolder)) {
        filesystem::create_directories(outImgFolder);
    }

    ofstream cameraOut(outFolder + "/camera.txt");
    ofstream imuOut(outFolder + "/imu.txt");
    cameraOut << fixed << setprecision(6);
    imuOut << fixed << setprecision(6);

    int idx1 = 0, idx2 = 0, idx3 = 0;
    while (imuVec[idx3 + 1].timestamp <= cameraTimeVec[idx1]) idx3++;
    while (videoCapture.isOpened() && idx1 < num_times) {
        Mat frame;
        string imageFileStr = outImgFolder + format("/%08d.png", idx1);
        if (videoCapture.read(frame)) {
            double t1 = cameraTimeVec[idx1];

            // sync gravity with camera
            while (gravityVec[idx2].timestamp < t1 && idx2 < num_gravity) idx2++;
            gravityVec[idx2].timestamp = t1;
            cameraOut << gravityVec[idx2] << endl;

            //sync imu with camera
            while (imuVec[idx3].timestamp < t1 && idx3 < num_imu) {
                imuOut << imuVec[idx3] << endl;
                idx3++;
            }

            cv::Mat image, grayImg;
            cv::undistort(frame, image, K, distort, rectK);
            cv::Rect rect(0, 0, rectW, rectH);
            image = image(rect);
            cv::cvtColor(image, grayImg, cv::COLOR_BGR2GRAY);
            imwrite(imageFileStr, grayImg);
            imshow("image", grayImg);
            waitKey(1);
            idx1++;
            cout << "iter: " << idx1 << " / " << num_times << endl;
        } else {
            cout << "Video end" << endl;
            break;
        }
    }

    if (idx3 < num_imu)
        imuOut << imuVec[idx3] << endl;

    videoCapture.release();
    cameraOut.close();
    imuOut.close();

    return 0;
}